package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"

	"github.com/duke-git/lancet/v2/system"
	"github.com/tebeka/selenium"
)

type CreateWebDriverOpts struct {
	Headless   bool // 静默模式
	DisableGpu bool // 禁用GPU
	NoSandbox  bool // 禁用沙盒模式
}

// getProcessUseMemory 获取进程使用内存
func getProcessUseMemory(processName string) (memoryUsage int, err error) {
	// 获取 Chrome 浏览器的进程 ID
	processID, err := getProcessID(processName)
	if err != nil {
		log.Fatalf("Error getting Chrome process ID: %v", err)
	}

	// 获取 Chrome 浏览器的内存使用情况
	memoryUsage, err = getProcessMemoryUsage(processID)
	if err != nil {
		return
	}

	return
}

// CreateWebDriver 创建浏览器会话
func CreateWebDriver(opts CreateWebDriverOpts) (driver selenium.WebDriver, err error) {
	// 配置浏览器的启动选项和功能
	caps := selenium.Capabilities{"browserName": "chrome"}

	chromeOptions := make([]string, 0)
	if opts.Headless {
		chromeOptions = append(chromeOptions, "--headless")
	}
	if opts.DisableGpu {
		chromeOptions = append(chromeOptions, "--disable-gpu")
	}
	if opts.NoSandbox {
		chromeOptions = append(chromeOptions, "--no-sandbox")
	}

	// Chrome配置
	var chromeCaps selenium.Capabilities
	chromeCaps = selenium.Capabilities{
		"goog:chromeOptions": map[string]interface{}{
			"args": chromeOptions,
		},
	}
	if system.IsLinux() {
		chromeCaps = selenium.Capabilities{
			"goog:chromeOptions": map[string]interface{}{
				// "binary": "./chrome/",
				"args": []string{
					"--headless", // 启用无头模式
					"--disable-gpu",
					"--no-sandbox",
					"--disable-dev-shm-usage",
					"--remote-debugging-port=9222", // 启用远程调试端口
				},
			},
		}
	}
	for k, v := range chromeCaps {
		caps[k] = v
	}

	// 创建新的WebDriver会话
	driver, err = selenium.NewRemote(caps, fmt.Sprintf("http://localhost:%d/wd/hub", port))
	if err != nil {
		log.Fatalf("无法创建浏览器会话: %v", err)
	}
	return
}

// startDriverService 启动驱动服务
func startDriverService() (service *selenium.Service, err error) {
	var chromeDriverPath string
	flag.StringVar(&chromeDriverPath, "d", "", "Chrome driver path")
	if chromeDriverPath == "" {
		extractChromeDriverPath, err_ := extractChromeDriver()
		if err_ != nil {
			err = errors.New(fmt.Sprintf("Selenium 驱动无法获取: %v", err_.Error()))
			return
		}
		chromeDriverPath = extractChromeDriverPath
	}

	// Selenium WebDriver 服务配置
	opts := make([]selenium.ServiceOption, 0)

	// Selenium WebDriver 服务启动
	service, err = selenium.NewChromeDriverService(chromeDriverPath, port, opts...)
	if err != nil {
		return
	}
	return
}

// extractChromeDriver 将 ChromeDriver 从嵌入的文件系统中提取到临时目录
func extractChromeDriver() (string, error) {
	var driverFileName string

	if system.IsMac() {
		driverFileName = fmt.Sprintf("embed/chromedriver_%s_%s", runtime.GOOS, runtime.GOARCH)
	}

	if system.IsWindows() {
		driverFileName = "embed/chromedriver.exe"
	}

	if system.IsLinux() {
		driverFileName = "embed/chromedriver"
	}

	data, err := embeddedFiles.ReadFile(driverFileName)
	if err != nil {
		return "", err
	}

	tmpFile, err := ioutil.TempFile("", "chromedriver-*")
	if err != nil {
		return "", err
	}
	defer tmpFile.Close()

	if _, err := tmpFile.Write(data); err != nil {
		return "", err
	}

	if err := os.Chmod(tmpFile.Name(), 0755); err != nil {
		return "", err
	}

	return tmpFile.Name(), nil
}

// getProcessID 获取进程 ID
func getProcessID(processName string) (int, error) {
	var cmd *exec.Cmd
	if runtime.GOOS == "windows" {
		// 使用 tasklist 和 findstr 组合查找包含 "chrome" 的进程
		cmd = exec.Command("cmd", "/C", fmt.Sprintf("tasklist | findstr %s", processName))
	} else {
		cmd = exec.Command("pgrep", "-f", processName)
	}

	output, err := cmd.Output()
	if err != nil {
		return 0, err
	}

	if runtime.GOOS == "windows" {
		lines := strings.Split(string(output), "\n")
		if len(lines) > 0 {
			columns := strings.Fields(lines[0])
			if len(columns) > 1 {
				return strconv.Atoi(columns[1])
			}
		}
		return 0, fmt.Errorf(fmt.Sprintf("no %s process found", processName))
	} else {
		processIDs := strings.Fields(string(output))
		if len(processIDs) == 0 {
			return 0, fmt.Errorf(fmt.Sprintf("no %s process found", processName))
		}
		return strconv.Atoi(processIDs[0])
	}
}

// getProcessMemoryUsage 获取指定进程的内存使用情况
func getProcessMemoryUsage(pid int) (int, error) {
	var cmd *exec.Cmd
	if runtime.GOOS == "windows" {
		cmd = exec.Command("tasklist", "/fi", fmt.Sprintf("PID eq %d", pid))
	} else {
		cmd = exec.Command("ps", "-p", strconv.Itoa(pid), "-o", "rss=")
	}

	output, err := cmd.Output()
	if err != nil {
		return 0, err
	}

	// 解析内存使用情况
	memoryUsage, err := strconv.Atoi(strings.TrimSpace(string(output)))
	if err != nil {
		return 0, err
	}

	return memoryUsage, nil
}
